/*
*    l3qd - Light, light, lightweight queue
*    Copyright (C) 2024  Marcus Pedersén marcus@marcux.org
*
*    This program is free software: you can redistribute it and/or modify
*    it under the terms of the GNU General Public License as published by
*    the Free Software Foundation, either version 3 of the License, or
*    (at your option) any later version.
*
*    This program is distributed in the hope that it will be useful,
*    but WITHOUT ANY WARRANTY; without even the implied warranty of
*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*    GNU General Public License for more details.
*
*    You should have received a copy of the GNU General Public License
*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package main

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/shirou/gopsutil/v4/cpu"
	"github.com/shirou/gopsutil/v4/mem"
	"github.com/spf13/viper"
	"gorm.io/gorm"
	"log/slog"
	"net"
	"net/http"
	"notabug.org/marcux/l3q/internal/common"
	"notabug.org/marcux/l3q/internal/common/cert"
	"notabug.org/marcux/l3q/internal/node"
	"notabug.org/marcux/l3q/internal/node/sql"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
)

const systemdUnitName = "l3qd-node-task-"

type Config struct {
	Port                int32  `mapstructure:"port"`
	IpListen            string `mapstructure:"ip_listen"`
	DbFile              string `mapstructure:"db_file"`
	LogFile             string `mapstructure:"log_file"`
	CertDir             string `mapstructure:"cert_directory"`
	L3qdPort            int32  `mapstructure:"l3qd_port"`
	SystemdSliceName    string `mapstructure:"systemd_slice_name"`
	StatusCheckInterval int    `mapstructure:"status_check_interval"`
	TaskOutputDirectory string `mapstructure:"task_output_directory"`
}

type API struct {
	Db   *gorm.DB
	Conf Config
}

// API for main daemon.
// URL: GET /api/v1/daemon/status
// Information about this node is returned.
func (a API) daemonGetStatus(c *gin.Context) {
	node := common.Node{}

	cpusInfo, err := cpu.Info()
	if err != nil {
		slog.Error(fmt.Sprintf("Request: /api/v1/daemon/status, failed to get CPU info, %s", err))
		c.JSON(http.StatusInternalServerError, gin.H{})
		return
	}

	for _, c := range cpusInfo {
		node.CpusTotal += int(c.Cores)
	}

	cpusPrcnt, err := cpu.Percent(0, true)
	if err != nil {
		slog.Error(fmt.Sprintf("Request: /api/v1/daemon/status, failed to get CPU percent usage, %s", err))
		c.JSON(http.StatusInternalServerError, gin.H{})
		return
	}

	for _, p := range cpusPrcnt {
		if p > 3.0 {
			node.CpusUsed++
		}
	}

	virtMem, err := mem.VirtualMemory()
	if err != nil {
		slog.Error(fmt.Sprintf("Request: /api/v1/daemon/status, failed to get memory usage, %s", err))
		c.JSON(http.StatusInternalServerError, gin.H{})
		return
	}

	swapMem, err := mem.SwapMemory()
	if err != nil {
		slog.Error(fmt.Sprintf("Request: /api/v1/daemon/status, failed to get swap usage, %s", err))
		c.JSON(http.StatusInternalServerError, gin.H{})
		return
	}

	node.MemoryTotal = virtMem.Total
	node.MemoryUsed = virtMem.Used
	node.SwapTotal = swapMem.Total
	node.SwapUsed = swapMem.Used

	c.JSON(http.StatusOK, node)
}

// API for main daemon.
// URL: POST /api/v1/daemon/task
// Posted tasks will be started
// by this node daemon
func (a API) daemonPostTask(c *gin.Context) {
	tasks := []common.NodeTask{}

	if err := c.Bind(&tasks); err != nil {
		slog.Error(fmt.Sprintf("Request: POST /api/v1/daemon/task, failed to parse tasks: %s", err))
		c.JSON(http.StatusBadRequest, common.FailResponse{Message: "Failed to parse tasks"})
		return
	}

	memoryHigh := "infinity"
	memoryMax := "infinity"

	errTasks := []common.NodeTask{}
	for _, t := range tasks {
		if t.Task.MemoryLimit > 0 {
			memoryHigh = fmt.Sprintf("%d", t.Task.MemoryLimit)
			memMax := float64(t.Task.MemoryLimit) * 1.05
			memoryMax = fmt.Sprintf("%d", uint64(memMax))
		}

		usr, err := user.Lookup(t.User)
		if err != nil {
			switch err.(type) {
			case user.UnknownUserError:
				slog.Error(fmt.Sprintf("Request: POST /api/v1/daemon/task, failed to lookup user: %s, User does not exist", t.User))
				t.Error = fmt.Sprintf("Failed to lookup user: %s, User does not exist", t.User)
			default:
				slog.Error(fmt.Sprintf("Request: POST /api/v1/daemon/task, failed to lookup user: %s, %s", t.User, err))
				t.Error = fmt.Sprintf("Failed to lookup user: %s", t.User)
			}

			errTasks = append(errTasks, t)
			continue
		}

		if t.Task.WorkDir == "$HOME" {
			t.Task.WorkDir = usr.HomeDir
		}

		grp, err := user.LookupGroup(t.Group)
		if err != nil {
			switch err.(type) {
			case user.UnknownGroupError:
				slog.Error(fmt.Sprintf("Request: POST /api/v1/daemon/task, failed to lookup group: %s, Group does not exist", t.Group))
				t.Error = fmt.Sprintf("Failed to lookup group: %s, Group does not exist", t.Group)
			default:
				slog.Error(fmt.Sprintf("Request: POST /api/v1/daemon/task, failed to lookup group: %s, %s", t.Group, err))
				t.Error = fmt.Sprintf("Failed to lookup group: %s", t.Group)
			}

			errTasks = append(errTasks, t)
			continue
		}

		cmdSystemd := exec.Command("/usr/bin/systemd-run",
			fmt.Sprintf("--property=AllowedCPUs=%d", t.Task.Cores),
			fmt.Sprintf("--property=MemoryHigh=%s", memoryHigh),
			fmt.Sprintf("--property=MemoryMax=%s", memoryMax),
			//systemd version >= 240
			fmt.Sprintf("--property=StandardOutput=append:%s",
				filepath.Join(a.Conf.TaskOutputDirectory,
					fmt.Sprintf("%s%d-%d.stdout", systemdUnitName, t.Task.JobID, t.Task.Id))),
			fmt.Sprintf("--property=StandardError=append:%s",
				filepath.Join(a.Conf.TaskOutputDirectory,
					fmt.Sprintf("%s%d-%d.stderr", systemdUnitName, t.Task.JobID, t.Task.Id))),
			fmt.Sprintf("--uid=%s", usr.Uid),
			fmt.Sprintf("--gid=%s", grp.Gid),
			"--no-ask-password",
			"--no-block",
			"--remain-after-exit",
			"--service-type=oneshot",
			fmt.Sprintf("--slice=%s", a.Conf.SystemdSliceName),
			fmt.Sprintf("--unit=%s%d-%d", systemdUnitName, t.Task.JobID, t.Task.Id),
			fmt.Sprintf("--working-directory=%s", t.Task.WorkDir),
			t.Task.Command)

		outBytes, err := cmdSystemd.CombinedOutput()
		if err != nil {
			slog.Error(fmt.Sprintf("Request: POST /api/v1/daemon/task, failed to execute systemd-run: %s, %s", string(outBytes), err))
			t.Error = "Failed to add task to systemd"
			errTasks = append(errTasks, t)
			continue
		}
	}

	if len(errTasks) > 0 {
		c.JSON(http.StatusInternalServerError, errTasks)
	} else {
		c.JSON(http.StatusOK, gin.H{})
	}
}

// API for main daemon.
// URL: PATCH /api/v1/daemon/task/cancel
// Cancel specified runing task running
// on this node.
// Parameter id and jobid is required.
//
// Query parameters:
// id=x: where x is task id that
// should be cancelled.
// jobid=y: where y is job id for task that
// should be cancelled.
func (a API) daemonPatchCancelTask(c *gin.Context) {
	idStr, iExist := c.GetQuery("id")
	if iExist {
		_, err := strconv.Atoi(idStr)
		if err != nil {
			slog.Error(fmt.Sprintf("Request: /api/v1/daemon/task/cancel?id=%s, id is not an integer", idStr))
			c.JSON(http.StatusBadRequest, common.FailResponse{Message: "Parameter id is not an integer"})
			return
		}
	} else {
		slog.Error("Request: /api/v1/daemon/task/cancel does not have required parameter: id")
		c.JSON(http.StatusBadRequest, common.FailResponse{Message: "Parameter id is required"})
		return
	}

	jobidStr, jExist := c.GetQuery("jobid")
	if jExist {
		_, err := strconv.Atoi(jobidStr)
		if err != nil {
			slog.Error(fmt.Sprintf("Request: /api/v1/daemon/task/cancel?jobid=%s, jobid is not an integer", jobidStr))
			c.JSON(http.StatusBadRequest, common.FailResponse{Message: "Parameter jobid is not an integer"})
			return
		}
	} else {
		slog.Error("Request: /api/v1/daemon/task/cancel does not have required parameter: jobid")
		c.JSON(http.StatusBadRequest, common.FailResponse{Message: "Parameter jobid is required"})
		return
	}

	unitName := fmt.Sprintf("%s%s-%s", systemdUnitName, idStr, jobidStr)

	taskBefore, err := node.StatusTask(a.Conf.SystemdSliceName, unitName, a.Conf.TaskOutputDirectory)
	if err != nil {
		slog.Error(fmt.Sprintf("Cancel task: Failed to get status before cancel for task: %s, job: %s, unit: %s, %s", idStr, jobidStr, unitName, err))
	}

	systemctlCmd := exec.Command("/usr/bin/systemctl", "kill", unitName)
	outBytes, err := systemctlCmd.CombinedOutput()
	if err != nil {
		slog.Error(fmt.Sprintf("Cancel task: Failed to run: systemctl kill %s, %s, %s", unitName, err, string(outBytes)))
		c.JSON(http.StatusInternalServerError, common.FailResponse{Message: "Failed to cancel task"})
		return
	}

	taskAfter, err := node.StatusTask(a.Conf.SystemdSliceName, unitName, a.Conf.TaskOutputDirectory)
	if err != nil {
		slog.Error(fmt.Sprintf("Cancel task: Failed to get status after cancel for task: %s, job: %s, unit: %s, %s", idStr, jobidStr, unitName, err))
	}

	if len(taskAfter.User) == 0 {
		taskAfter.User = taskBefore.User
	}

	if len(taskAfter.Group) == 0 {
		taskAfter.Group = taskBefore.Group
	}

	if taskAfter.Task.Id == 0 {
		taskAfter.Task.Id = taskBefore.Task.Id
	}

	if taskAfter.Task.JobID == 0 {
		taskAfter.Task.JobID = taskBefore.Task.JobID
	}

	if len(taskAfter.Task.Unitname) == 0 {
		taskAfter.Task.Unitname = taskBefore.Task.Unitname
	}

	if taskAfter.Task.MainPid == 0 {
		taskAfter.Task.MainPid = taskBefore.Task.MainPid
	}

	if len(taskAfter.Task.StartTime) == 0 {
		taskAfter.Task.StartTime = taskBefore.Task.StartTime
	}

	if taskAfter.Task.MonoStartTime == 0 {
		taskAfter.Task.MonoStartTime = taskBefore.Task.MonoStartTime
	}

	if len(taskAfter.Task.EndTime) == 0 {
		taskAfter.Task.EndTime = taskBefore.Task.EndTime
	}

	if taskAfter.Task.MonoEndTime == 0 {
		taskAfter.Task.MonoEndTime = taskBefore.Task.MonoEndTime
	}

	if len(taskAfter.Task.ActiveState) == 0 {
		taskAfter.Task.ActiveState = taskBefore.Task.ActiveState
	}

	if len(taskAfter.Task.SubState) == 0 {
		taskAfter.Task.SubState = taskBefore.Task.SubState
	}

	if taskAfter.Task.ExitCode == 0 {
		taskAfter.Task.ExitCode = taskBefore.Task.ExitCode
	}

	if taskAfter.Task.MemoryPeak == 0 {
		taskAfter.Task.MemoryPeak = taskBefore.Task.MemoryPeak
	}

	if taskAfter.Task.PidsPeak == 0 {
		taskAfter.Task.PidsPeak = taskBefore.Task.PidsPeak
	}

	if len(taskAfter.Task.Stdout) == 0 {
		taskAfter.Task.Stdout = taskBefore.Task.Stdout
	}

	if len(taskAfter.Task.Stderr) == 0 {
		taskAfter.Task.Stderr = taskBefore.Task.Stderr
	}

	if taskAfter.Task.SubState == "failed" {
		sysFailCmd := exec.Command("/usr/bin/systemctl", "reset-failed", unitName)
		outBytes, err := sysFailCmd.CombinedOutput()
		if err != nil {
			slog.Error(fmt.Sprintf("Cancel task: Failed to run: systemctl reser-failed after cancel for task: %s, job: %s, unit: %s, output: %s, %s",
				idStr, jobidStr, unitName, string(outBytes), err))
			taskAfter.Task.ActiveState = "Node-Err"
		}
	} else {
		// systemctl stop "unitname"
		// if fail
		// taskAfter.ActiveState = "Node-Err"
		sysStopCmd := exec.Command("/usr/bin/systemctl", "stop", unitName)
		outBytes, err := sysStopCmd.CombinedOutput()
		if err != nil {
			slog.Error(fmt.Sprintf("Cancel task: Failed to run: systemctl stop after cancel for task: %s, job: %s, unit: %s, output: %s, %s",
				idStr, jobidStr, unitName, string(outBytes), err))
			taskAfter.Task.ActiveState = "Node-Err"
		}
	}

	c.JSON(http.StatusOK, taskAfter)
}

// Called if URL is invalid
func wrongUrl(c *gin.Context) {
	c.JSON(http.StatusNotFound, gin.H{})
}

// Read and parse config file
func getConfig() Config {
	exPath, err := os.Executable()
	if err != nil {
		slog.Error(fmt.Sprintf("Failed to lookup path to this executable, in trying to get config, %s", err))
		os.Exit(1)
	}

	viper.SetConfigName("l3qd-node.conf")
	viper.SetConfigType("toml")
	viper.AddConfigPath(".")
	viper.AddConfigPath("/etc/l3qd-node/")
	viper.AddConfigPath(filepath.Dir(exPath))
	viper.SetDefault("port", 39911)
	viper.SetDefault("ip_listen", "127.0.0.1")
	viper.SetDefault("db_file", "/var/lib/l3q/l3qd-node.db")
	viper.SetDefault("log_file", "/var/log/l3q/l3qd-node.log")
	viper.SetDefault("l3qd_port", 39911)
	viper.SetDefault("systemd_slice_name", "l3q.slice")
	viper.SetDefault("status_check_interval", 60)
	viper.SetDefault("task_output_directory", "/var/lib/l3q/l3qd-node-task-output/")
	viper.SetDefault("cert_directory", "/etc/l3qd-node/certificates")
	err = viper.ReadInConfig()
	if err != nil {
		slog.Error(fmt.Sprintf("Failed to read config file: l3qd-node.conf, %s", err))
		os.Exit(1)
	}

	var conf Config

	err = viper.Unmarshal(&conf)
	if err != nil {
		slog.Error(fmt.Sprintf("Failed to parse config file: l3qd-node.conf, %s", err))
		os.Exit(2)
	}

	if conf.Port <= 0 || conf.Port > 65535 {
		slog.Error(fmt.Sprintf("Config file: l3qd-node.conf, port: %d  is not a valid port number.", conf.Port))
		os.Exit(3)
	}

	if conf.L3qdPort <= 0 || conf.L3qdPort > 65535 {
		slog.Error(fmt.Sprintf("Config file: l3qd-node.conf, l3qd_port: %d  is not a valid port number.", conf.L3qdPort))
		os.Exit(3)
	}

	if net.ParseIP(conf.IpListen) == nil {
		slog.Error(fmt.Sprintf("Config file: l3qd-node.conf, IP addres: %s is not a valid IP address.", conf.IpListen))
		os.Exit(3)
	}

	if conf.StatusCheckInterval < 0 {
		slog.Error(fmt.Sprintf("Config file: l3qd-node.conf, status_check_interval: %d has to be a postiv integer.", conf.StatusCheckInterval))
		os.Exit(3)
	}

	// 86400 is 24 hours
	if conf.StatusCheckInterval > 86400 {
		slog.Warn(fmt.Sprintf("Config file: l3qd-node.conf, status_check_interval: %d is a large integer, larger then 24 hours", conf.StatusCheckInterval))
	}

	if len(conf.SystemdSliceName) == 0 {
		slog.Error("Config file: l3qd-node.conf, systemd_slice_name: can not be an empty string")
		os.Exit(3)
	}

	if len(conf.TaskOutputDirectory) == 0 {
		slog.Error("Config file: l3qd-node.conf, task_output_directory: can not be an empty string")
		os.Exit(3)
	}

	if _, err = os.Stat(conf.TaskOutputDirectory); err != nil {
		if os.IsNotExist(err) {
			slog.Error(fmt.Sprintf("Config file: l3qd-node.conf, task_output_directory: %s does not exist.", conf.TaskOutputDirectory))
			os.Exit(3)
		} else {
			slog.Warn(fmt.Sprintf("Config file: l3qd-node.conf, Failed to access task_output_directory: %s.", conf.TaskOutputDirectory))
		}
	}

	if len(conf.DbFile) == 0 {
		slog.Error("Config file: l3qd-node.conf, db_file: can not be an empty string")
		os.Exit(3)
	}

	if _, err = os.Stat(filepath.Dir(conf.DbFile)); err != nil {
		if os.IsNotExist(err) {
			slog.Error(fmt.Sprintf("Config file: l3qd-node.conf, db_file directory: %s does not exist.", filepath.Dir(conf.DbFile)))
			os.Exit(3)
		} else {
			slog.Warn(fmt.Sprintf("Config file: l3qd-node.conf, Failed to access db_file directory: %s.", filepath.Dir(conf.DbFile)))
		}
	}

	if len(conf.LogFile) == 0 {
		slog.Error("Config file: l3qd-node.conf, log_file: can not be an empty string")
		os.Exit(3)
	}

	if _, err = os.Stat(filepath.Dir(conf.LogFile)); err != nil {
		if os.IsNotExist(err) {
			slog.Error(fmt.Sprintf("Config file: l3qd-node.conf, log_file directory: %s does not exist.", filepath.Dir(conf.LogFile)))
			os.Exit(3)
		} else {
			slog.Warn(fmt.Sprintf("Config file: l3qd-node.conf, Failed to access log_file directory: %s.", filepath.Dir(conf.LogFile)))
		}
	}

	return conf
}

// Validate this node to the main l3q daemon.
func validate() error {
	// TODO
	fmt.Println("Validating ...")
	return nil
}

// Prints all running tasks to stdout
func printStatus() {
	conf := getConfig()
	tasks, err := node.AllTaskStatus(conf.SystemdSliceName, "")
	if err != nil {
		slog.Error(fmt.Sprintf("%s", err))
		fmt.Fprintf(os.Stderr, "Failed to get status for all tasks. Output is incomplete.")
	}

	fmt.Println(strings.Repeat("-", 153))
	fmt.Printf("| %5s | %6s |    %4s    |          %10s          |           %8s           | %4s |   %5s   |                 %7s                |\n",
		"Jobid", "Taskid", "User", "Start time", "End time", "Exit", "State", "Systemd")
	fmt.Println(strings.Repeat("-", 153))

	for _, t := range tasks {
		starttime, endtime, exit := "", "", ""

		if len(t.Task.StartTime) == 0 {
			starttime = "-"
		} else {
			starttime = t.Task.StartTime
		}

		if len(t.Task.EndTime) == 0 {
			endtime = "-"
		} else {
			endtime = t.Task.EndTime
		}

		if endtime == "-" {
			exit = "-"
		} else {
			exit = fmt.Sprintf("%d", t.Task.ExitCode)
		}

		fmt.Printf("| %5d | %6d | %10s | %28s | %28s | %4s | %9s | %39s |\n",
			t.Task.JobID, t.Task.Id, t.User, starttime, endtime, exit, t.Task.SubState, t.Task.Unitname)
		fmt.Println(strings.Repeat("-", 153))
	}
}

// Prints version
func printVersion() {
	common.AsciiL3q()
	fmt.Printf("l3qd %s\n", common.VERSION)
	fmt.Printf("%s\n", common.COPYRIGHT)
	fmt.Println("License GPLv3+: GNU GPL version 3 or later <https://gnu.org/licenses/gpl.html>.")
	fmt.Println("This is free software: you are free to change and redistribute it.")
	fmt.Println("There is NO WARRANTY, to the extent permitted by law.")
	fmt.Println("")
	fmt.Println("Written by Marcus Pedersén.")
}

// Prints help text
func printHelp() {
	fmt.Println("Usage: l3qd-node [OPTION]...")
	common.AsciiL3q()
	fmt.Println("The L3q daemon that runs on each compute node.")
	fmt.Println("l3qd-node communicates with central L3q daemon (l3qd)")
	fmt.Println("and execute tasks on this compute node.")
	fmt.Println("")
	fmt.Println("Run without arguments (daemon mode):")
	fmt.Println("Will start in daemon mode, start processing")
	fmt.Println("tasks and communicating with central L3q")
	fmt.Println("daemon (l3qd). This is the normal operation")
	fmt.Println("mode. If daemon is not running status on")
	fmt.Println("running tasks will not be updated.")
	fmt.Println("If daemon is not running new tasks can not")
	fmt.Println("be received from l3qd and new tasks will not")
	fmt.Println("be launched. Central l3q daemon (l3qd) will")
	fmt.Println("show this node as offline as it can not be reached.")
	fmt.Println("")
	fmt.Println("Options:")
	fmt.Println("    --status         Print information for tasks on this node.")
	fmt.Println("    --validate       Without validation the node daemon on this host will")
	fmt.Println("                     not be able to communicate with L3q daemon.")
	fmt.Println("                     Start the validation server on the same server")
	fmt.Println("                     as the L3q daemon with:")
	fmt.Println("                     l3qd --validate-host")
	fmt.Println("                     and follow instructions.")
	fmt.Println("                     root previlegies are required to run this command.")
	fmt.Println("    --version,       Prints version and exit.")
	fmt.Println("    --help, -h       Prints this help text and exit.")
	fmt.Println("")
	fmt.Println("Documentation available locally via: man l3qd-node")
	fmt.Println("Report bugs to: <https://notabug.org/marcux/l3q/issues>")
}

func main() {
	if len(os.Args) == 2 {
		switch os.Args[1] {
		case "--status":
			printStatus()
			os.Exit(0)
		case "--version":
			printVersion()
			os.Exit(0)
		case "--help", "-h":
			printHelp()
			os.Exit(0)
		case "--validate":
			if err := validate(); err != nil {
				os.Exit(1)
			} else {
				os.Exit(0)
			}
		default:
			fmt.Fprintf(os.Stderr, "Wrong arguments. Try -h or --help for help.\n")
			os.Exit(1)
		}
	} else if len(os.Args) > 2 {
		fmt.Fprintf(os.Stderr, "Wrong arguments. Try -h or --help for help.\n")
		os.Exit(1)
	}

	conf := getConfig()

	db, err := sql.InitDb(conf.DbFile)
	if err != nil {
		slog.Error(fmt.Sprintf("%s", err))
		os.Exit(2)
	}

	api := API{}
	api.Conf = conf
	api.Db = db

	if err := cert.GenerateCert(conf.CertDir); err != nil {
		slog.Error(fmt.Sprintf("%s", err))
		os.Exit(3)
	}

	route := gin.Default()
	route.NoRoute(wrongUrl)
	route.NoMethod(wrongUrl)
	route.GET("/api/v1/daemon/status", api.daemonGetStatus)
	route.POST("/api/v1/daemon/task", api.daemonPostTask)
	route.PATCH("/api/v1/daemon/task/cancel", api.daemonPatchCancelTask)

	route.RunTLS(fmt.Sprintf("%s:%d", conf.IpListen, conf.Port),
		filepath.Join(conf.CertDir, "cert.pem"),
		filepath.Join(conf.CertDir, "key.pem"))
}
